#include "ImageProcessor.h"
#include <QDebug>
#include <QThread>

ImageProcessor::ImageProcessor(QObject *parent)
    : QObject(parent), m_codec(nullptr), m_codecCtx(nullptr),
      m_frame(nullptr), m_packet(nullptr), m_frameCount(0)
{
    // 初始化编码器
    if (!initEncoder(320, 256, 30)) {
        emit errorOccurred("编码器初始化失败");
    }
}

ImageProcessor::~ImageProcessor()
{
    QMutexLocker locker(&m_mutex);
    
    // 发送空帧，确保编码器缓存中的数据被刷新
    if (m_codecCtx) {
        avcodec_send_frame(m_codecCtx, nullptr);
    }
    
    // 释放FFmpeg资源
    if (m_packet) {
        av_packet_free(&m_packet);
        m_packet = nullptr;
    }
    if (m_frame) {
        av_frame_free(&m_frame);
        m_frame = nullptr;
    }
    if (m_codecCtx) {
        avcodec_free_context(&m_codecCtx);
        m_codecCtx = nullptr;
    }
}

bool ImageProcessor::initEncoder(int width, int height, int fps)
{
    QMutexLocker locker(&m_mutex);
    
    // 查找H.264编码器
    m_codec = avcodec_find_encoder(AV_CODEC_ID_H264);
    if (!m_codec) {
        emit errorOccurred("找不到H.264编码器");
        return false;
    }
    
    // 创建编码器上下文
    m_codecCtx = avcodec_alloc_context3(m_codec);
    if (!m_codecCtx) {
        emit errorOccurred("无法分配编码器上下文");
        return false;
    }
    
    // 设置编码器参数
    m_codecCtx->width = width;
    m_codecCtx->height = height;
    m_codecCtx->time_base = {1, fps};
    m_codecCtx->framerate = {fps, 1};
    m_codecCtx->gop_size = 10; // 每10帧一个关键帧
    m_codecCtx->max_b_frames = 1;
    m_codecCtx->pix_fmt = AV_PIX_FMT_YUV420P;
    m_codecCtx->bit_rate = 2000000; // 2Mbps比特率
    
    // 对于实时传输，使用零延迟模式
    if (m_codec->id == AV_CODEC_ID_H264) {
        av_opt_set(m_codecCtx->priv_data, "preset", "ultrafast", 0);
        av_opt_set(m_codecCtx->priv_data, "tune", "zerolatency", 0);
    }
    
    // 打开编码器
    if (avcodec_open2(m_codecCtx, m_codec, nullptr) < 0) {
        emit errorOccurred("无法打开编码器");
        avcodec_free_context(&m_codecCtx);
        m_codecCtx = nullptr;
        return false;
    }
    
    // 初始化帧和包
    m_frame = av_frame_alloc();
    if (!m_frame) {
        emit errorOccurred("无法分配帧");
        avcodec_free_context(&m_codecCtx);
        m_codecCtx = nullptr;
        return false;
    }
    m_frame->format = m_codecCtx->pix_fmt;
    m_frame->width = width;
    m_frame->height = height;
    
    // 分配帧数据缓冲区
    if (av_frame_get_buffer(m_frame, 0) < 0) {
        emit errorOccurred("无法为帧分配缓冲区");
        av_frame_free(&m_frame);
        avcodec_free_context(&m_codecCtx);
        m_frame = nullptr;
        m_codecCtx = nullptr;
        return false;
    }
    
    m_packet = av_packet_alloc();
    if (!m_packet) {
        emit errorOccurred("无法分配数据包");
        av_frame_free(&m_frame);
        avcodec_free_context(&m_codecCtx);
        m_frame = nullptr;
        m_codecCtx = nullptr;
        return false;
    }
    
    qInfo() << "编码器初始化成功，分辨率:" << width << "x" << height << "帧率:" << fps;
    return true;
}

bool ImageProcessor::convertToYUV420P(const QByteArray &grayData, AVFrame *frame)
{
    // 16位灰度图转8位并填充到YUV420P的Y分量
    const uint16_t *src = reinterpret_cast<const uint16_t*>(grayData.constData());
    uint8_t *yPlane = frame->data[0];
    int yStride = frame->linesize[0];
    
    for (int y = 0; y < frame->height; y++) {
        for (int x = 0; x < frame->width; x++) {
            // 16位转8位：右移8位或根据实际范围调整
            uint8_t val = static_cast<uint8_t>(src[y * frame->width + x] >> 8);
            yPlane[y * yStride + x] = val;
        }
    }
    
    // 填充UV分量为128（灰度图，UV分量为中间值）
    uint8_t *uPlane = frame->data[1];
    uint8_t *vPlane = frame->data[2];
    int uStride = frame->linesize[1];
    int vStride = frame->linesize[2];
    int uvHeight = frame->height / 2;
    int uvWidth = frame->width / 2;
    
    for (int y = 0; y < uvHeight; y++) {
        memset(uPlane + y * uStride, 128, uvWidth);
        memset(vPlane + y * vStride, 128, uvWidth);
    }
    
    return true;
}

bool ImageProcessor::encodeFrame(AVFrame *frame, QByteArray &outputData)
{
    QMutexLocker locker(&m_mutex);
    if (!m_codecCtx || !frame || !m_packet) {
        return false;
    }
    
    // 设置帧序号和时间戳
    frame->pts = m_frameCount++;
    
    // 发送帧到编码器
    int ret = avcodec_send_frame(m_codecCtx, frame);
    if (ret < 0) {
        emit errorOccurred("发送帧到编码器失败");
        return false;
    }
    
    // 接收编码后的数据包
    ret = avcodec_receive_packet(m_codecCtx, m_packet);
    if (ret == AVERROR(EAGAIN) || ret == AVERROR_EOF) {
        return false; // 暂时没有输出，或已完成
    } else if (ret < 0) {
        emit errorOccurred("编码帧失败");
        return false;
    }
    
    // 将编码数据存入QByteArray
    outputData.append(reinterpret_cast<const char*>(m_packet->data), m_packet->size);
    
    // 释放数据包数据（但保留数据包结构）
    av_packet_unref(m_packet);
    return true;
}

void ImageProcessor::processImageData(const QByteArray &data)
{
    // 检查是否在工作线程中运行
    if (QThread::currentThread() == QCoreApplication::instance()->thread()) {
        emit errorOccurred("图像处理不在工作线程中运行");
        return;
    }
    
    QMutexLocker locker(&m_mutex);
    if (!m_frame || !m_codecCtx) {
        emit errorOccurred("编码器未初始化");
        return;
    }
    
    // 确保帧数据可写
    if (av_frame_make_writable(m_frame) < 0) {
        emit errorOccurred("无法获取帧的写入权限");
        return;
    }
    
    // 转换16位灰度图到YUV420P格式
    if (!convertToYUV420P(data, m_frame)) {
        emit errorOccurred("转换图像格式失败");
        return;
    }
    
    // 编码帧
    QByteArray encodedData;
    if (encodeFrame(m_frame, encodedData) && !encodedData.isEmpty()) {
        emit frameEncoded(encodedData);
    }
}
